import random
import pygame

# Initialize Pygame
pygame.init()

# Set up the window
WIDTH = 800
HEIGHT = 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Lunar Lander")

# Color definitions
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)


# Lander class
class Lander:
    def __init__(self, fuel=500, thrust=0.2, gravity=0.1):
        self.x = WIDTH // 2
        self.y = 50
        self.velocity_y = 0
        self.velocity_x = 0
        self.fuel = fuel
        self.thrust = thrust
        self.gravity = gravity
        self.size = 20

    def update(self, wind=0):
        # Update position
        self.velocity_y += self.gravity
        self.velocity_x += wind  # Apply wind effect
        self.y += self.velocity_y
        self.x += self.velocity_x

        # Boundary check
        if self.x < 0:
            self.x = WIDTH
        elif self.x > WIDTH:
            self.x = 0

    def thrust_up(self):
        if self.fuel > 0:
            self.velocity_y -= self.thrust
            self.fuel -= 0.5

    def thrust_left(self):
        if self.fuel > 0:
            self.velocity_x -= self.thrust / 2
            self.fuel -= 0.3

    def thrust_right(self):
        if self.fuel > 0:
            self.velocity_x += self.thrust / 2
            self.fuel -= 0.3

    def draw(self):
        pygame.draw.polygon(screen, WHITE, [
            (self.x, self.y),
            (self.x - self.size / 2, self.y + self.size),
            (self.x + self.size / 2, self.y + self.size)
        ])


# Create terrain
def create_terrain():
    terrain = []
    y = HEIGHT - 100
    for x in range(0, WIDTH + 10, 10):
        y += random.randint(-5, 5)
        terrain.append((x, y))
    return terrain


# Input screen for game parameters
def input_screen():
    font = pygame.font.Font(None, 36)
    small_font = pygame.font.Font(None, 28)

    # Default values
    fuel = "500"
    gravity = "0.1"
    thrust = "0.2"

    active_field = None

    while True:
        screen.fill(BLACK)

        # Title
        title = font.render("LUNAR LANDER - SETTINGS", True, WHITE)
        screen.blit(title, (WIDTH // 2 - title.get_width() // 2, 50))

        # Input fields
        fuel_text = font.render(f"Fuel: {fuel}", True, WHITE if active_field != 'fuel' else GREEN)
        gravity_text = font.render(f"Gravity: {gravity}", True, WHITE if active_field != 'gravity' else GREEN)
        thrust_text = font.render(f"Thrust: {thrust}", True, WHITE if active_field != 'thrust' else GREEN)

        screen.blit(fuel_text, (100, 150))
        screen.blit(gravity_text, (100, 200))
        screen.blit(thrust_text, (100, 250))

        # Instructions
        instructions = [
            "Click on a field to edit it",
            "Press ENTER to start the game",
            "Use UP/DOWN/LEFT/RIGHT to control the lander",
            "Land gently on the green platform to win"
        ]

        for i, instruction in enumerate(instructions):
            text = small_font.render(instruction, True, WHITE)
            screen.blit(text, (100, 350 + i * 30))

        pygame.display.flip()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return None, None, None

            if event.type == pygame.MOUSEBUTTONDOWN:
                # Check which field was clicked
                mouse_x, mouse_y = pygame.mouse.get_pos()
                if 100 <= mouse_x <= 300 and 150 <= mouse_y <= 180:
                    active_field = 'fuel'
                elif 100 <= mouse_x <= 300 and 200 <= mouse_y <= 230:
                    active_field = 'gravity'
                elif 100 <= mouse_x <= 300 and 250 <= mouse_y <= 280:
                    active_field = 'thrust'
                else:
                    active_field = None

            if event.type == pygame.KEYDOWN:
                if active_field:
                    if event.key == pygame.K_BACKSPACE:
                        if active_field == 'fuel':
                            fuel = fuel[:-1]
                        elif active_field == 'gravity':
                            gravity = gravity[:-1]
                        elif active_field == 'thrust':
                            thrust = thrust[:-1]
                    elif event.key == pygame.K_RETURN:
                        # Try to convert inputs to numbers and start game
                        try:
                            fuel_val = float(fuel)
                            gravity_val = float(gravity)
                            thrust_val = float(thrust)
                            return fuel_val, gravity_val, thrust_val
                        except ValueError:
                            pass  # Invalid input, stay on input screen
                    else:
                        # Add character to active field
                        char = event.unicode
                        if char.isdigit() or char == '.':
                            if active_field == 'fuel':
                                fuel += char
                            elif active_field == 'gravity':
                                gravity += char
                            elif active_field == 'thrust':
                                thrust += char
                elif event.key == pygame.K_RETURN:
                    # Try to convert inputs to numbers and start game
                    try:
                        fuel_val = float(fuel)
                        gravity_val = float(gravity)
                        thrust_val = float(thrust)
                        return fuel_val, gravity_val, thrust_val
                    except ValueError:
                        pass  # Invalid input, stay on input screen


# Main game loop
def main():
    def reset_game(fuel, gravity, thrust):
        # Reset game state
        nonlocal lander, terrain, landing_pad, game_over, success, wind
        lander = Lander(fuel, thrust, gravity)
        terrain = create_terrain()
        landing_pad = random.randint(WIDTH // 4, WIDTH * 3 // 4)
        game_over = False
        success = False
        wind = (random.random() - 0.5) * 0.05  # Random wind effect

    clock = pygame.time.Clock()

    # Show input screen and get parameters
    fuel, gravity, thrust = input_screen()
    if fuel is None:  # User closed the window
        return

    lander = Lander(fuel, thrust, gravity)
    terrain = create_terrain()
    landing_pad = random.randint(WIDTH // 4, WIDTH * 3 // 4)
    game_over = False
    success = False

    # Wind effect - random constant force
    wind = (random.random() - 0.5) * 0.05

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return
            # Add restart game functionality with R key
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r:
                    reset_game(fuel, gravity, thrust)

        if not game_over:
            # Controls
            keys = pygame.key.get_pressed()
            if keys[pygame.K_UP]:
                lander.thrust_up()
            if keys[pygame.K_LEFT]:
                lander.thrust_left()
            if keys[pygame.K_RIGHT]:
                lander.thrust_right()

            lander.update(wind)  # Pass wind to update method

            # Collision detection
            for i in range(len(terrain) - 1):
                if terrain[i][0] <= lander.x <= terrain[i + 1][0]:
                    ground_y = terrain[i][1]
                    if lander.y + lander.size >= ground_y:
                        game_over = True
                        # Check for safe landing
                        if (abs(lander.velocity_y) < 2 and
                                abs(lander.velocity_x) < 2 and
                                abs(lander.x - landing_pad) < 30):
                            success = True

        # Drawing
        screen.fill(BLACK)

        # Draw terrain
        pygame.draw.lines(screen, WHITE, False, terrain)

        # Draw landing pad
        pygame.draw.line(screen, GREEN,
                         (landing_pad - 20, terrain[landing_pad // 10][1]),
                         (landing_pad + 20, terrain[landing_pad // 10][1]), 3)

        lander.draw()

        # Display fuel
        font = pygame.font.Font(None, 36)
        fuel_text = font.render(f'Fuel: {int(lander.fuel)}', True, WHITE)
        screen.blit(fuel_text, (10, 10))

        # Display speed
        velocity_text = font.render(f'Speed: {abs(lander.velocity_y):.1f}', True, WHITE)
        screen.blit(velocity_text, (10, 50))

        # Display wind
        wind_text = font.render(f'Wind: {wind:.3f}', True, BLUE if wind > 0 else WHITE)
        screen.blit(wind_text, (10, 90))

        # Display game result and restart prompt
        if game_over:
            result_text = font.render('SUCCESS!' if success else 'CRASH!', True,
                                      GREEN if success else RED)
            restart_text = font.render('Press R to restart', True, WHITE)
            screen.blit(result_text, (WIDTH // 2 - 50, HEIGHT // 2))
            screen.blit(restart_text, (WIDTH // 2 - 70, HEIGHT // 2 + 40))

        pygame.display.flip()
        clock.tick(60)


if __name__ == '__main__':
    main()
